import joblib
from sklearn.datasets._base import Bunch
from sklearn.feature_extraction.text import TfidfVectorizer
from Tools import readfile, readbunchobj, writebunchobj

def vector_space(stopword_path, bunch_path, space_path, train_tfidf_path=None):

    stopword_list = readfile(stopword_path).splitlines()
    set_bunch = readbunchobj(bunch_path)

    tfidf_space = Bunch(target_name=set_bunch.target_name, label=set_bunch.label, filenames=set_bunch.filenames, tdm=[], vocabulary={})

    if train_tfidf_path is None: # 训练集处理
        vectorizer = TfidfVectorizer(stop_words=stopword_list, max_features=40000, sublinear_tf=True, max_df=0.5)
        '''
        #.stop_words=list类型,直接过滤指定的停用词。
        # sublinear_tf:，计算tf值采用亚线性策略。比如，我们以前算tf是词频，现在用1+log(tf)来充当词频。
        # max_df,过滤出现在超过max_df=0.5比例的句子中的词语,当他在全文档出现的频次过多>50%时我们认为他太过常见而不具备代表性
        '''
        tfidf_space.tdm = vectorizer.fit_transform(set_bunch.contents)
        tfidf_space.vocabulary = vectorizer.vocabulary_

    else: # 测试集处理 有train_tfidf_path这个参数
        trainbunch = readbunchobj(train_tfidf_path)  # 读取训练集的TF-IDF模型
        tfidf_space.vocabulary = trainbunch.vocabulary  # 设置词汇表,测试集将使用与训练集相同的词汇表

        # 创建TF-IDF向量化器
        '''
            # stop_words=list类型,直接过滤指定的停用词。
            # sublinear_tf:，计算tf值采用亚线性策略。比如，我们以前算tf是词频，现在用1+log(tf)来充当词频,有助于减少高频词的权重
            # max_df,过滤出现在超过max_df=0.5比例的句子中的词语,当他在全文档出现的频次过多>50%时我们认为他太过常见而不具备代表性,被视为过于常见而被忽略。
            # vocabulary: dict类型,只使用特定的词汇(训练集中的)，为了避免在测试集中出现训练集中没有出现的词汇而造成困扰所以一般会用这个，但是如果训练集足够大可以不用
        '''
        vectorizer = TfidfVectorizer(stop_words=stopword_list, max_features=4000, sublinear_tf=True, max_df=0.5,
                                     vocabulary=trainbunch.vocabulary)

        # 转换测试集文本
        '''
            # fit_transform方法用于拟合模型并转换文本数据。这里，它使用训练集的词汇表来转换测试集bunch.contents中的文本
            # 结果是一个TF-IDF矩阵（术语-文档矩阵），存储在tfidf_space.tdm中
            # 行对应输入文本，列对应词汇表中的词汇，矩阵元素值表示相应词汇在文本中以TF-IDF权重表示的重要性。
        '''
        tfidf_space.tdm = vectorizer.fit_transform(set_bunch.contents)

    # 保存vectorizer对象
    joblib.dump(vectorizer, 'vectorizer.joblib')

    writebunchobj(space_path, tfidf_space)
    print("if-idf词向量空间实例创建成功！")


if __name__ == '__main__':

    stopword_path = "stopwords.txt"

    # 训练集
    bunch_path = "train_word_bag/train_set.dat"
    space_path = "train_word_bag/tf_idf_space.dat"
    vector_space(stopword_path, bunch_path, space_path)

    # 测试集
    train_tfidf_path = "train_word_bag/tf_idf_space.dat"
    bunch_path = "test_word_bag/test_set.dat"
    space_path = "test_word_bag/test_tf_idf_space.dat"
    vector_space(stopword_path, bunch_path, space_path, train_tfidf_path)
